
import os, sys
from mscf.fc3.project.volfitter.gui.comparewindow import CompareDialog

os.environ['ETS_TOOLKIT'] = 'qt4'
from PySide.QtGui import QMainWindow
from pyface.qt import QtGui, QtCore

import numpy
from traits.api import HasTraits, Instance, on_trait_change
from traitsui.api import View, Item

from mayavi.core.ui.api import MayaviScene, MlabSceneModel,SceneEditor

import matplotlib
matplotlib.use('Qt4Agg')
matplotlib.rcParams['backend.qt4']='PySide'
import matplotlib.pyplot as plt
from matplotlib.backends.backend_qt4agg import FigureCanvasQTAgg as FigureCanvas

from mpl_toolkits.mplot3d import axes3d

from pylab import *

from mscf.fc3.project.volfitter.marketdata.optionmarketdata import  OptionMarketData
from mscf.fc3.project.volfitter.analytics.volsurface import  ParametricVolSurface
from mscf.fc3.project.volfitter.analytics.volsurface import  DTM
from mscf.fc3.project.volfitter.analytics.volsurfacefitter import  VolSurfaceFitter

from mscf.fc3.project.volfitter.analytics.optionpricing import TheoPricer

from mscf.fc3.project.volfitter.gui.utils import GUIHelper

import datetime
from dateutil.relativedelta import relativedelta

from mscf.fc3.project.volfitter.marketdata.marketdataaccess import MarketDataFetcher


APP_NAME = "Vol Fitter"

class MainApp(QMainWindow):

    def __init__(self,  guiData, parent=None):
        super(MainApp, self).__init__(parent)
        self.guiData = guiData
        self.fitVolSurface()
        self.frame = MainFrame(guiData)
        self.setWindowTitle(APP_NAME)
        self.setCentralWidget(self.frame)
        screen = QtGui.QDesktopWidget().screenGeometry()
#        self.move((screen.width()  - self.__WND_WIDTH)  / 2,  (screen.height() - self.__WND_HEIGHT) / 2)
#        size = self.sizeHint()
#        size.setWidth(screen.width() * 3 /4 )
#        size.setHeight(screen.height() * 3 / 4)
        self.setGeometry(100, 150, screen.width() * 3 /4 , screen.height() * 3 / 4)
        del screen # Small Cleanup
        #self.showMaximized()

    def fitVolSurface(self):
        volSurfaceFitter = VolSurfaceFitter(gsOptionMarketData)
        self.guiData.currentVolSurface = volSurfaceFitter.fitParametricVolSurface()
        #self.displayVSPlots()

class GUIData():

    def setCurrentTicker(self, currentTicker):
        self.currentTicker = currentTicker

    def setCurrentMarketData(self, optionMarketData):
        self.currentOptionMarketData = optionMarketData

    def setCurrentExpDate(self, currentExpDate):
        self.currentExpDate = currentExpDate

    def setCurrentVolSurface(self, currentVolSurface):
        self.currentVolSurface = currentVolSurface





class MainFrame(QtGui.QDialog):

    optionDataHeaderLabels = ['strike','Call Bid', 'Call Theo', 'Call Ask', 'Impl Vol','Put Bid', 'Put Theo', 'Put Ask']


    def __init__(self,  guiData):
        super(MainFrame, self).__init__()

        #self.setGeometry(0,0, 500,650)
        self.setWindowIcon(QtGui.QIcon("vsfitter.jpg"))
        self.center()

        self.guiData = guiData

        self.createMenu()

        outerLayout = QtGui.QVBoxLayout(self)

        controlPanel = self.createControlPanel()
        outputPanel = self.createOutputPanel()

        mainPanel = QtGui.QSplitter(QtCore.Qt.Vertical)
        mainPanel.addWidget(controlPanel)
        mainPanel.addWidget(outputPanel)

        outerLayout.addWidget(mainPanel)
        outerLayout.setMenuBar(self.menuBar)
        self.setLayout(outerLayout)



    def createControlPanel(self):
        controlPanel = QtGui.QGroupBox("Ticker Control")

        controlLayout = QtGui.QHBoxLayout()
        controlPanel.setLayout(controlLayout)

        tickerLabel = QtGui.QLabel("Ticker")
        controlLayout.addWidget(tickerLabel)

        self.tickerBox = QtGui.QLineEdit()
        controlLayout.addWidget(self.tickerBox)

        mktDataButton = QtGui.QPushButton("Get Option Prices")
        mktDataButton.clicked.connect(self.displayMarketData)
        controlLayout.addWidget(mktDataButton)

        fitbutton = QtGui.QPushButton("Fit Volatility Surface")
        fitbutton.clicked.connect(self.fitVolSurface)
        controlLayout.addWidget(fitbutton)

        return controlPanel

    def createOutputPanel(self):
        outputPanel = QtGui.QSplitter(QtCore.Qt.Horizontal)

        currentVolCurve2DFig = GUIHelper.createVolCurveFit2DFigure(self.guiData.currentVolSurface, self.guiData.currentExpDate, self.guiData.currentOptionMarketData)
        currentVolSurf3DFig = GUIHelper.createVolSurf3DFigure(self.guiData.currentVolSurface)
        self.plotViewPanel = QtGui.QSplitter(QtCore.Qt.Vertical)
        self.plotViewPanel.addWidget(FigureCanvas(currentVolCurve2DFig))
        self.plotViewPanel.addWidget(currentVolSurf3DFig.canvas)

        self.dataPanel = self.createDataPanel()
        self.dataPanel.setFrameShape(QtGui.QFrame.StyledPanel)
        outputPanel.addWidget(self.dataPanel)
        outputPanel.addWidget(self.plotViewPanel)

        w = outputPanel.width()
        outputPanel.setSizes ([3*w/5, 2*w/5])

        return outputPanel


    def createDataPanel(self):
        dataPanel = QtGui.QFrame(self)
        dataPanel.setFrameShape(QtGui.QFrame.StyledPanel)

        return dataPanel


    def displayMarketData(self):

        optionMarketData = self.guiData.currentOptionMarketData

        self.optionDataTabWidget = QtGui.QTabWidget()

        for expDate, optionPriceChain in sorted(optionMarketData.optionPriceChainMap.iteritems()):



            callPutPairMap = optionPriceChain.callPutPairMap
            optionDataTable = QtGui.QTableWidget(len(optionPriceChain.callPutPairMap), len(self.optionDataHeaderLabels))
            optionDataTable.setHorizontalHeaderLabels(self.optionDataHeaderLabels)
            optionDataTable.horizontalHeader().setResizeMode(QtGui.QHeaderView.Stretch)
            self.optionDataTabWidget.addTab(optionDataTable, expDate.isoformat())

            i = 0
            for k, v in sorted(callPutPairMap.iteritems()):
                cell = QtGui.QTableWidgetItem()
                cell.setText(str(k))
                optionDataTable.setItem(i, 0,  cell)
                if not(v is None):

                    if not(v.callPrice is None):
                        price = v.callPrice
                        if not(price.bidPrice is None):
                            cell = QtGui.QTableWidgetItem()
                            cell.setText(str(price.bidPrice))
                            optionDataTable.setItem(i, 1, cell)

                        if ( expDate == self.guiData.currentExpDate ):
                            spot = optionMarketData.refSpot
                            strike = price.strike
                            interestRate = optionMarketData.interestRate
                            term = self.guiData.currentVolSurface.dtmMap[expDate]
                            vol = self.guiData.currentVolSurface.lookupVolatility(strike, term)
                            theo = TheoPricer.BlsPrice(spot ,strike, interestRate, term, vol, "call")

                            cell = QtGui.QTableWidgetItem()
                            cell.setText(str(theo))
                            optionDataTable.setItem(i, 2, cell)


                        if not(price.askPrice is None):
                            cell = QtGui.QTableWidgetItem()
                            cell.setText(str(price.askPrice))
                            optionDataTable.setItem(i, 3, cell)


                        if not(price.impVol is None):
                            cell = QtGui.QTableWidgetItem()
                            cell.setText(str(price.impVol))
                            optionDataTable.setItem(i, 4, cell)

                    if not(v.putPrice is None):
                        price = v.putPrice
                        if not(price.bidPrice is None):
                            cell = QtGui.QTableWidgetItem()
                            cell.setText(str(price.bidPrice))
                            optionDataTable.setItem(i, 5, cell)

                        if ( expDate == self.guiData.currentExpDate ):
                            spot = optionMarketData.refSpot
                            strike = price.strike
                            interestRate = optionMarketData.interestRate
                            term = self.guiData.currentVolSurface.dtmMap[expDate]
                            vol = self.guiData.currentVolSurface.lookupVolatility(strike, term)
                            theo = TheoPricer.BlsPrice(spot ,strike, interestRate, term, vol, "put")

                            cell = QtGui.QTableWidgetItem()
                            cell.setText(str(theo))
                            optionDataTable.setItem(i, 6, cell)

                        if not(price.askPrice is None):
                            cell = QtGui.QTableWidgetItem()
                            cell.setText(str(price.askPrice))
                            optionDataTable.setItem(i, 7, cell)

                    i += 1


        dataPanelLayout = QtGui.QVBoxLayout()
        dataPanelLayout.addWidget(self.optionDataTabWidget)

        self.dataPanel.setLayout(dataPanelLayout)



    def displayMarketData2(self):

        theoPricer = TheoPricer()

        optionMarketData = gsOptionMarketData

        optionDataTabWidget = self.optionDataTabWidget

        optionDataTabWidget.ta

        for expDate, optionPriceChain in sorted(optionMarketData.optionPriceChainMap.iteritems()):

            callPutPairMap = optionPriceChain.callPutPairMap
            optionDataTable = QtGui.QTableWidget(len(optionPriceChain.callPutPairMap), len(self.optionDataHeaderLabels))
            optionDataTable.setHorizontalHeaderLabels(self.optionDataHeaderLabels)
            optionDataTable.horizontalHeader().setResizeMode(QtGui.QHeaderView.Stretch)
            optionDataTabWidget.addTab(optionDataTable, expDate.isoformat())

            i = 0
            for k, v in sorted(callPutPairMap.iteritems()):
                cell = QtGui.QTableWidgetItem()
                cell.setText(str(k))
                optionDataTable.setItem(i, 0,  cell)
                if not(v is None):

                    if not(v.callPrice is None):
                        price = v.callPrice
                        if not(price.bidPrice is None):
                            cell = QtGui.QTableWidgetItem()
                            cell.setText(str(price.bidPrice))
                            optionDataTable.setItem(i, 1, cell)


                        spot = optionMarketData.refSpot
                        strike = price.strike
                        interestRate = optionMarketData.interestRate
                        term = self.guiData.currentVolSurface.dtmMap[expDate]
                        vol = self.guiData.currentVolSurface.lookupVolatility(strike, term)
                        theo = TheoPricer.BlsPrice(spot ,strike, interestRate, term, vol, "call")

                        cell = QtGui.QTableWidgetItem()
                        cell.setText(str(theo))
                        optionDataTable.setItem(i, 2, cell)


                        if not(price.askPrice is None):
                            cell = QtGui.QTableWidgetItem()
                            cell.setText(str(price.askPrice))
                            optionDataTable.setItem(i, 3, cell)
                        if not(price.impVol is None):
                            cell = QtGui.QTableWidgetItem()
                            cell.setText(str(price.impVol))
                            optionDataTable.setItem(i, 4, cell)

                    if not(v.putPrice is None):
                        price = v.putPrice
                        if not(price.bidPrice is None):
                            cell = QtGui.QTableWidgetItem()
                            cell.setText(str(price.bidPrice))
                            optionDataTable.setItem(i, 5, cell)
                        if not(price.askPrice is None):
                            cell = QtGui.QTableWidgetItem()
                            cell.setText(str(price.askPrice))
                            optionDataTable.setItem(i, 7, cell)

                    i += 1


        dataPanelLayout = QtGui.QVBoxLayout()
        dataPanelLayout.addWidget(optionDataTabWidget)

        self.dataPanel.setLayout(dataPanelLayout)


    def displayVSPlots(self):
        currentVolCurve2DFig = GUIHelper.createVolCurveFit2DFigure(self.guiData.currentVolSurface, self.guiData.currentExpDate, self.guiData.currentOptionMarketData)
        currentVolSurf3DFig = GUIHelper.createVolSurf3DFigure(self.guiData.currentVolSurface)
#        self.plotViewPanel = QtGui.QSplitter(QtCore.Qt.Vertical)
        self.plotViewPanel.addWidget(FigureCanvas(currentVolCurve2DFig))
        self.plotViewPanel.addWidget(currentVolSurf3DFig.canvas)

    def fitVolSurface(self):
        volSurfaceFitter = VolSurfaceFitter(109, 0.02, self.guiData.currentOptionMarketData)
        self.guiData.currentVolSurface = volSurfaceFitter.fitParametricVolSurface()
        self.displayVSPlots()

    def createMenu(self):
        self.menuBar = QtGui.QMenuBar()

        self.fileMenu = QtGui.QMenu(self.tr("&File"), self)
        self.compareAction = self.fileMenu.addAction(self.tr("&Compare"))
        self.exitAction = self.fileMenu.addAction(self.tr("E&xit"))
        self.menuBar.addMenu(self.fileMenu)
        self.compareAction.triggered.connect(self.handleCompare)

        self.connect(self.exitAction, QtCore.SIGNAL("triggered()"),self.accept)

    def handleCompare(self):
        self.compareWindow = CompareDialog(self.guiData.currentVolSurface)
        self.compareWindow.show()


    def center(self):
        screen = QtGui.QDesktopWidget().screenGeometry()
        size = self.geometry()
        self.move((screen.width()-size.width())/2, (screen.height()-size.height())/2)





testTicker = "IBM"
refSpot = 203.00
ir = 0.05
#testOptionMarketData = MarketDataFetcher.fetchMarketData(testTicker, refSpot)
#testOptionMarketData.interestRate=ir


app = QtCore.QCoreApplication.instance()

testExpDate = datetime.date(2012,10,19)
#datetime.datetime.now() + relativedelta( years = +1 )
#dtms=[]
#sixMonthsMinus = testExpDate + relativedelta( months = -6 )
#dtms.append(DTM(sixMonthsMinus, 0.5))
#dtms.append(DTM(testExpDate, 1.0))
#sixMonthsPlus = testExpDate + relativedelta( months = +6 )
#dtms.append(DTM(sixMonthsPlus, 1.5))
#testVolSurface = ParametricVolSurface(1200.0, None, 0.05, [0.5, 1.0, 1.5], dtms)


refSpot = 109
gsOptionMarketData = OptionMarketData.getTestGSData(refSpot)
gsOptionMarketData.interestRate = ir
gsOptionMarketData2 = OptionMarketData.getTestGSData2(refSpot)
gsOptionMarketData2.interestRate = ir
#getTestData(testExpDate)

testGUIData = GUIData()
testGUIData.setCurrentTicker(testTicker)
#testGUIData.setCurrentVolSurface(testVolSurface)
testGUIData.setCurrentExpDate(testExpDate)
testGUIData.setCurrentMarketData(gsOptionMarketData2)

main = MainApp(testGUIData)
main.show()

sys.exit(app.exec_())